Skip to content

Preserve tracked array wrappers in restructure - #498

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
JuliaArrays:masterfrom
ChrisRackauckas-Claude:agent/preserve-tracked-restructure
Draft

Preserve tracked array wrappers in restructure#498
ChrisRackauckas-Claude wants to merge 1 commit into
JuliaArrays:masterfrom
ChrisRackauckas-Claude:agent/preserve-tracked-restructure

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown

Ignore this PR until it has been reviewed by @ChrisRackauckas.

Summary

  • Preserve ReverseDiff and Tracker tracked-array wrappers when the restructure template is tracked.
  • Convert arrays of tracked scalars back to their tracked array representation before reshaping.
  • Cover tracked-array and array-of-tracked-scalars inputs in the AD integration tests.

Why

ArrayInterface.restructure(tracked_template, tracked_values) currently falls through to the generic implementation. For an array of ReverseDiff.TrackedReal values, that produces a plain vector, so a downstream Rosenbrock stage changes from ReverseDiff.TrackedArray to Vector{TrackedReal} and cannot be assigned to the typed stage vector.

This was exposed by the SciMLSensitivity Core1 BouncingBall ODE test on clean current masters after the OrdinaryDiffEq LinearSolve 5 update. The fix belongs in the public ArrayInterface.restructure extensions rather than in an OrdinaryDiffEq workaround.

Investigation process

  1. Reproduced the clean-master failure under ReverseDiff, Tracker, and Zygote outer backends.
  2. Isolated the failing value as Vector{ReverseDiff.TrackedReal{..., Nothing}} and confirmed the generic restructure stripped the wrapper.
  3. Rejected a downstream _restructure_state fallback after it failed the exact solve.
  4. Verified that aos_to_soa plus reshape preserves the tracked wrapper and gradient, then implemented that behavior in the owner extensions.
  5. Ran the configured JuliaFormatter on the touched files.

Local validation

  • GROUP=Core julia +1.10 --project=. -e "using Pkg; Pkg.test()"
    • BandedMatrices 21/21
    • BlockBandedMatrices 8/8
    • Core 221/221
    • AD Integration 28/28
    • StaticArrays 39/39
    • ChainRules 20/20
    • FillArrays 7/7
    • process result: Testing ArrayInterface tests passed
  • Focused stage probe preserved the wrapper through restructure, scalar multiplication, and typed stage assignment; ReverseDiff gradient [2.0, 2.0].
  • Unchanged Rosenbrock23 plus SciMLSensitivity BouncingBall reproduction returned the stable repeated ReverseDiff gradient [26.96250030687455, 0.0].
  • Original ForwardDiff/FiniteDiff reference plus ReverseDiff, Tracker, and Zygote matrix: 4/4 passed in 38.2 seconds.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@codecov

codecov Bot commented Jul 29, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 61.32%. Comparing base (d767284) to head (b3e2a79).

Additional details and impacted files
@@            Coverage Diff             @@
##           master     #498      +/-   ##
==========================================
+ Coverage   60.75%   61.32%   +0.57%     
==========================================
  Files          15       15              
  Lines         609      618       +9     
==========================================
+ Hits          370      379       +9     
  Misses        239      239              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants